import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import time

def train_model(model, train_loader, criterion, optimizer, scheduler, epochs=50, device=None):
    """
    训练模型
    
    参数:
        model (nn.Module): 要训练的模型
        train_loader (DataLoader): 训练数据加载器
        criterion (nn.Module): 损失函数
        optimizer (optim.Optimizer): 优化器
        scheduler (optim.lr_scheduler): 学习率调度器
        epochs (int): 训练轮数
        device (torch.device): 训练设备(CPU/GPU)
        
    返回:
        model (nn.Module): 训练好的模型
        train_losses (list): 每轮的训练损失
        train_accuracies (list): 每轮的训练准确率
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    model.to(device)
    model.train()
    
    train_losses = []
    train_accuracies = []
    
    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        start_time = time.time()
        
        # 使用tqdm显示进度条
        loop = tqdm(train_loader, leave=False)
        for batch_idx, (inputs, labels) in enumerate(loop):
            inputs, labels = inputs.to(device), labels.to(device)
            
            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # 统计信息
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            
            # 更新进度条
            loop.set_description(f'Epoch [{epoch+1}/{epochs}]')
            loop.set_postfix(loss=loss.item(), acc=100.*correct/total)
        
        # 更新学习率
        scheduler.step()
        
        # 计算本轮平均损失和准确率
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100. * correct / total
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc)
        
        # 打印训练信息
        epoch_time = time.time() - start_time
        print(f'Epoch {epoch+1}/{epochs} - Loss: {epoch_loss:.4f} - Acc: {epoch_acc:.2f}% - Time: {epoch_time:.2f}s - LR: {scheduler.get_last_lr()[0]:.6f}')
    
    return model, train_losses, train_accuracies